"""
任务状态查询API
提供任务状态查询、列表、取消等功能
"""
from fastapi import APIRouter, HTTPException, Query, Depends
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from datetime import datetime, timedelta
from enum import Enum
import sys
import os

# 添加task-queue路径
task_queue_path = os.getenv('TASK_QUEUE_PATH',
                            os.path.abspath(os.path.join(
                                os.path.dirname(__file__),
                                '../../../../task-queue/src'
                            )))
sys.path.insert(0, task_queue_path)

from celery.result import AsyncResult
from celery import states
from src.celery_app import app as celery_app
import motor.motor_asyncio

# 导入认证模块
from .auth import AuthUser, get_current_user, require_permission, Permissions

router = APIRouter(prefix="/api/v1/tasks", tags=["tasks"])

# MongoDB客户端
mongo_client = motor.motor_asyncio.AsyncIOMotorClient(
    os.getenv('MONGODB_URI', 'mongodb://localhost:27017')
)
db = mongo_client[os.getenv('MONGODB_DB', 'crawler_scheduler')]


class TaskState(str, Enum):
    """任务状态枚举"""
    PENDING = "PENDING"
    STARTED = "STARTED"
    RETRY = "RETRY"
    SUCCESS = "SUCCESS"
    FAILURE = "FAILURE"
    REVOKED = "REVOKED"
    PROGRESS = "PROGRESS"


class TaskStatus(BaseModel):
    """任务状态模型"""
    task_id: str = Field(..., description="任务ID")
    state: TaskState = Field(..., description="任务状态")
    current: Optional[int] = Field(None, description="当前进度")
    total: Optional[int] = Field(None, description="总进度")
    result: Optional[Dict[str, Any]] = Field(None, description="任务结果")
    error: Optional[str] = Field(None, description="错误信息")
    meta: Optional[Dict[str, Any]] = Field(None, description="元数据")
    started_at: Optional[datetime] = Field(None, description="开始时间")
    completed_at: Optional[datetime] = Field(None, description="完成时间")
    duration: Optional[float] = Field(None, description="执行时长（秒）")


class TaskListResponse(BaseModel):
    """任务列表响应模型"""
    tasks: List[TaskStatus] = Field(..., description="任务列表")
    total: int = Field(..., description="总数")
    page: int = Field(..., description="当前页")
    per_page: int = Field(..., description="每页数量")
    has_next: bool = Field(..., description="是否有下一页")


class TaskStatistics(BaseModel):
    """任务统计模型"""
    total_tasks: int = Field(..., description="总任务数")
    pending: int = Field(..., description="待处理任务数")
    running: int = Field(..., description="运行中任务数")
    success: int = Field(..., description="成功任务数")
    failed: int = Field(..., description="失败任务数")
    success_rate: float = Field(..., description="成功率")
    avg_duration: float = Field(..., description="平均执行时长（秒）")
    tasks_per_hour: float = Field(..., description="每小时任务数")


@router.get("/{task_id}/status", response_model=TaskStatus)
async def get_task_status(
    task_id: str,
    current_user: AuthUser = Depends(get_current_user)
):
    """
    获取任务状态
    
    Args:
        task_id: 任务ID
        current_user: 当前认证用户
        
    Returns:
        任务状态信息
    """
    # 获取Celery任务状态
    async_result = AsyncResult(task_id, app=celery_app)
    
    # 查询数据库中的任务记录
    collection = db.task_results
    task_doc = await collection.find_one({'task_id': task_id})
    
    # 构建响应
    response = TaskStatus(
        task_id=task_id,
        state=TaskState(async_result.state)
    )
    
    # 根据状态添加详细信息
    if async_result.state == states.PENDING:
        response.meta = {"message": "Task is waiting to be executed"}
    
    elif async_result.state == states.STARTED:
        response.started_at = task_doc['started_at'] if task_doc else None
        response.meta = {"message": "Task is currently running"}
    
    elif async_result.state == 'PROGRESS':
        response.current = async_result.info.get('current', 0)
        response.total = async_result.info.get('total', 100)
        response.meta = async_result.info.get('meta', {})
    
    elif async_result.state == states.SUCCESS:
        response.result = async_result.result
        if task_doc:
            response.started_at = task_doc.get('started_at')
            response.completed_at = task_doc.get('completed_at')
            if response.started_at and response.completed_at:
                response.duration = (
                    response.completed_at - response.started_at
                ).total_seconds()
    
    elif async_result.state == states.FAILURE:
        response.error = str(async_result.info)
        if task_doc:
            response.started_at = task_doc.get('started_at')
            response.completed_at = task_doc.get('failed_at')
    
    elif async_result.state == states.REVOKED:
        response.meta = {"message": "Task was cancelled"}
    
    return response


@router.get("/", response_model=TaskListResponse)
async def list_tasks(
    page: int = Query(1, ge=1, description="页码"),
    per_page: int = Query(20, ge=1, le=100, description="每页数量"),
    state: Optional[TaskState] = Query(None, description="筛选状态"),
    platform: Optional[str] = Query(None, description="筛选平台"),
    start_time: Optional[datetime] = Query(None, description="开始时间"),
    end_time: Optional[datetime] = Query(None, description="结束时间"),
    current_user: AuthUser = Depends(get_current_user)
):
    """
    获取任务列表
    
    Args:
        page: 页码
        per_page: 每页数量
        state: 筛选状态
        platform: 筛选平台
        start_time: 开始时间
        end_time: 结束时间
        current_user: 当前认证用户
        
    Returns:
        任务列表
    """
    # 构建查询条件
    query = {}
    
    if state:
        query['state'] = state.value
    
    if platform:
        query['platform'] = platform
    
    if start_time or end_time:
        query['created_at'] = {}
        if start_time:
            query['created_at']['$gte'] = start_time
        if end_time:
            query['created_at']['$lte'] = end_time
    
    # 查询数据库
    collection = db.task_results
    
    # 获取总数
    total = await collection.count_documents(query)
    
    # 分页查询
    skip = (page - 1) * per_page
    cursor = collection.find(query).skip(skip).limit(per_page).sort('created_at', -1)
    
    tasks = []
    async for doc in cursor:
        # 获取Celery状态
        async_result = AsyncResult(doc['task_id'], app=celery_app)
        
        task_status = TaskStatus(
            task_id=doc['task_id'],
            state=TaskState(async_result.state),
            started_at=doc.get('started_at'),
            completed_at=doc.get('completed_at'),
            result=doc.get('result'),
            error=doc.get('error'),
            meta={
                'platform': doc.get('platform'),
                'triggered_by': doc.get('triggered_by')
            }
        )
        
        if task_status.started_at and task_status.completed_at:
            task_status.duration = (
                task_status.completed_at - task_status.started_at
            ).total_seconds()
        
        tasks.append(task_status)
    
    return TaskListResponse(
        tasks=tasks,
        total=total,
        page=page,
        per_page=per_page,
        has_next=(page * per_page < total)
    )


@router.delete("/{task_id}")
async def cancel_task(
    task_id: str,
    current_user: AuthUser = Depends(require_permission(Permissions.TASK_CANCEL))
):
    """
    取消任务（需要特定权限）
    
    Args:
        task_id: 任务ID
        current_user: 当前认证用户
        
    Returns:
        取消结果
    """
    try:
        # 取消Celery任务
        celery_app.control.revoke(task_id, terminate=True)
        
        # 更新数据库状态
        collection = db.task_results
        await collection.update_one(
            {'task_id': task_id},
            {
                '$set': {
                    'state': 'REVOKED',
                    'cancelled_at': datetime.utcnow(),
                    'cancelled_by': current_user.id
                }
            }
        )
        
        return {
            'success': True,
            'task_id': task_id,
            'message': 'Task cancelled successfully'
        }
        
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Failed to cancel task: {str(e)}"
        )


@router.get("/{task_id}/result")
async def get_task_result(
    task_id: str,
    current_user: AuthUser = Depends(get_current_user)
):
    """
    获取任务结果
    
    Args:
        task_id: 任务ID
        current_user: 当前认证用户
        
    Returns:
        任务结果
    """
    # 获取Celery任务结果
    async_result = AsyncResult(task_id, app=celery_app)
    
    if async_result.state != states.SUCCESS:
        raise HTTPException(
            status_code=400,
            detail=f"Task is not completed. Current state: {async_result.state}"
        )
    
    return {
        'task_id': task_id,
        'state': async_result.state,
        'result': async_result.result,
        'traceback': async_result.traceback
    }


@router.get("/statistics", response_model=TaskStatistics)
async def get_task_statistics(
    hours: int = Query(24, ge=1, le=168, description="统计时间范围（小时）"),
    platform: Optional[str] = Query(None, description="筛选平台"),
    current_user: AuthUser = Depends(get_current_user)
):
    """
    获取任务统计信息
    
    Args:
        hours: 统计时间范围（小时）
        platform: 筛选平台
        current_user: 当前认证用户
        
    Returns:
        任务统计信息
    """
    # 计算时间范围
    end_time = datetime.utcnow()
    start_time = end_time - timedelta(hours=hours)
    
    # 构建查询条件
    query = {
        'created_at': {
            '$gte': start_time,
            '$lte': end_time
        }
    }
    
    if platform:
        query['platform'] = platform
    
    # 查询数据库
    collection = db.task_results
    
    # 使用聚合管道计算统计
    pipeline = [
        {'$match': query},
        {
            '$group': {
                '_id': None,
                'total': {'$sum': 1},
                'pending': {
                    '$sum': {'$cond': [{'$eq': ['$state', 'PENDING']}, 1, 0]}
                },
                'running': {
                    '$sum': {'$cond': [
                        {'$in': ['$state', ['STARTED', 'RETRY', 'PROGRESS']]}, 
                        1, 0
                    ]}
                },
                'success': {
                    '$sum': {'$cond': [{'$eq': ['$state', 'SUCCESS']}, 1, 0]}
                },
                'failed': {
                    '$sum': {'$cond': [{'$eq': ['$state', 'FAILURE']}, 1, 0]}
                },
                'total_duration': {
                    '$sum': {
                        '$cond': [
                            {'$and': [
                                {'$ne': ['$duration', None]},
                                {'$eq': ['$state', 'SUCCESS']}
                            ]},
                            '$duration',
                            0
                        ]
                    }
                }
            }
        }
    ]
    
    cursor = collection.aggregate(pipeline)
    stats = await cursor.to_list(1)
    
    if not stats:
        return TaskStatistics(
            total_tasks=0,
            pending=0,
            running=0,
            success=0,
            failed=0,
            success_rate=0.0,
            avg_duration=0.0,
            tasks_per_hour=0.0
        )
    
    stats = stats[0]
    
    success_rate = (
        stats['success'] / stats['total'] 
        if stats['total'] > 0 else 0.0
    )
    
    avg_duration = (
        stats['total_duration'] / stats['success']
        if stats['success'] > 0 else 0.0
    )
    
    tasks_per_hour = stats['total'] / hours if hours > 0 else 0.0
    
    return TaskStatistics(
        total_tasks=stats['total'],
        pending=stats['pending'],
        running=stats['running'],
        success=stats['success'],
        failed=stats['failed'],
        success_rate=round(success_rate, 2),
        avg_duration=round(avg_duration, 2),
        tasks_per_hour=round(tasks_per_hour, 2)
    )


@router.websocket("/ws/{task_id}")
async def websocket_task_updates(
    websocket,
    task_id: str
):
    """
    WebSocket实时任务状态更新
    
    Args:
        websocket: WebSocket连接
        task_id: 任务ID
    """
    await websocket.accept()
    
    try:
        while True:
            # 获取任务状态
            async_result = AsyncResult(task_id, app=celery_app)
            
            # 发送状态更新
            await websocket.send_json({
                'task_id': task_id,
                'state': async_result.state,
                'info': async_result.info if async_result.info else None
            })
            
            # 如果任务完成，关闭连接
            if async_result.state in [states.SUCCESS, states.FAILURE, states.REVOKED]:
                break
            
            # 等待1秒再更新
            import asyncio
            await asyncio.sleep(1)
            
    except Exception as e:
        await websocket.send_json({
            'error': str(e)
        })
    finally:
        await websocket.close()